package ai.example.langchain4j.controller;

import ai.example.langchain4j.config.SseEmitter;
import ai.example.langchain4j.service.Assistant;
import ai.example.langchain4j.service.SteamAssistant;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import lombok.RequiredArgsConstructor;
import org.springframework.http.MediaType;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.util.List;

import static ai.example.langchain4j.utils.ChatMessageBuilder.buildChatMessage;


@RestController
@CrossOrigin
@RequiredArgsConstructor
@RequestMapping("/api")
public class LangChainChatController {

  private final ChatModel chatModel;
  private final Assistant assistant;
  private final SteamAssistant steamAssistant;
  private final ChatMemoryProvider chatMemoryProvider;
  private final StreamingChatModel streamingChatModel;

  @PostMapping("/chat")
  public String chat(@RequestParam(required = false) String id,
                     @RequestParam String query,
                     @RequestParam(value = "file", required = false) List<MultipartFile> files) {

    ChatMessage userMessage = buildChatMessage(query, files);

    if (StringUtils.hasLength(id)) {
      ChatMemory chatMemory = chatMemoryProvider.get(id);
      chatMemory.add(userMessage);
      String result = chatModel.chat(chatMemory.messages()).aiMessage().text();
      chatMemory.add(new AiMessage(result));
      return result;
    }

    return chatModel.chat(userMessage).aiMessage().text();
  }

  @PostMapping(value = "/chat/sse", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
  public Flux<String> steamAssistant(@RequestParam(required = false) String id,
                                     @RequestParam String query,
                                     @RequestParam(value = "file", required = false) List<MultipartFile> files) {
    ChatMemory chatMemory = StringUtils.hasLength(id) ? chatMemoryProvider.get(id) : null;
    ChatMessage userMessage = buildChatMessage(query, files);
    return Flux.create(new SseEmitter(chatMemory, userMessage, streamingChatModel));
  }


  @GetMapping("/chat/simplified")
  public String assistant3(@RequestParam String id,
                           @RequestParam String query) {
    String result = assistant.chat(id, query);
    return result;
  }


}
